"""
进程管理器 - 负责跟踪和关闭所有子进程和异步任务
基于注册表管理，不进行系统范围的进程检索
"""

import os
import signal
import psutil
import threading
import time
import logging
import asyncio
from typing import List, Dict, Set, Optional, Awaitable, Any
from pathlib import Path
import atexit
from dataclasses import dataclass, field
from contextlib import contextmanager, asynccontextmanager
from concurrent.futures import ThreadPoolExecutor, as_completed
import weakref
import uuid
from datetime import datetime

# 导入配置
from config import Config

@dataclass
class ProcessInfo:
    """进程信息"""
    pid: int
    name: str
    process_type: str
    start_time: float
    task_id: Optional[str] = None
    status: str = field(default="running")
    
    def __post_init__(self):
        """检查进程是否仍然存在"""
        try:
            psutil.Process(self.pid)
        except psutil.NoSuchProcess:
            self.status = "terminated"
    
    def is_alive(self) -> bool:
        """检查进程是否仍然存活"""
        try:
            proc = psutil.Process(self.pid)
            return proc.is_running()
        except psutil.NoSuchProcess:
            self.status = "terminated"
            return False

@dataclass
class AsyncTaskInfo:
    """异步任务信息"""
    task_id: str
    name: str
    task_type: str  # 'edge_connection', 'batch_operation', 'sample_collection'
    start_time: float
    task: Optional[asyncio.Task] = None
    status: str = field(default="running")
    related_nodes: Optional[List[str]] = None
    
    def is_done(self) -> bool:
        """检查任务是否已完成"""
        if self.task is None:
            return True
        return self.task.done()
    
    def cancel(self) -> bool:
        """取消任务"""
        if self.task and not self.task.done():
            self.task.cancel()
            self.status = "cancelled"
            return True
        return False

class ProcessManager:
    """进程管理器 - 线程安全的单例模式，基于注册表管理进程和异步任务"""
    
    _instance = None
    _lock = threading.RLock()  # 使用可重入锁
    
    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialize()
        return cls._instance
    
    def _initialize(self):
        """初始化方法，只调用一次"""
        if hasattr(self, '_initialized'):
            return
        
        self._initialized = True
        self.logger = logging.getLogger(__name__)
        
        # 使用锁保护的数据结构
        self._data_lock = threading.RLock()
        
        # 主要的进程注册表 - 只管理这里注册的进程
        self._registered_processes: Dict[int, ProcessInfo] = {}
        self._process_by_task: Dict[str, Set[int]] = {}  # 按任务ID分组的进程
        
        # 异步任务注册表
        self._async_tasks: Dict[str, AsyncTaskInfo] = {}
        self._async_tasks_by_type: Dict[str, Set[str]] = {}  # 按类型分组的异步任务
        
        # 训练线程跟踪
        self._training_threads: Dict[str, threading.Thread] = {}
        self._active_tasks: Set[str] = set()
        
        # 使用弱引用避免循环引用
        self._cleanup_callbacks = weakref.WeakSet()
        
        # 信号处理标志
        self._shutdown_initiated = threading.Event()
        self._shutdown_timeout = 30  # 30秒超时
        
        # 线程池用于并发清理
        self._cleanup_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="cleanup")
        
        # 注册退出处理
        atexit.register(self.cleanup_all)
        
        # 注册信号处理器
        self._setup_signal_handlers()
        
        self.logger.info("进程管理器初始化完成 - 基于注册表模式，支持异步任务管理")
    
    def _setup_signal_handlers(self):
        """设置信号处理器"""
        def signal_handler(signum, frame):
            self.logger.info(f"接收到信号 {signum}，开始关闭...")
            self.graceful_shutdown()
        
        # 注册常见的关闭信号
        for sig in [signal.SIGTERM, signal.SIGINT]:
            try:
                signal.signal(sig, signal_handler)
            except (ValueError, OSError):
                # 某些信号可能无法在当前环境中注册
                self.logger.warning(f"无法注册信号 {sig}")
    
    def _safe_log(self, message: str, level: str = "info"):
        """安全的日志记录方法，避免在日志系统关闭后记录日志"""
        try:
            # 检查logger是否仍然可用
            if hasattr(self, 'logger') and self.logger and hasattr(self.logger, 'handlers'):
                # 检查是否有可用的处理器
                if self.logger.handlers:
                    # 检查处理器是否仍然可用
                    for handler in self.logger.handlers:
                        if hasattr(handler, 'stream') and not handler.stream.closed:
                            if level == "error":
                                self.logger.error(message)
                            elif level == "warning":
                                self.logger.warning(message)
                            else:
                                self.logger.info(message)
                            return
            # 如果日志系统不可用，静默忽略
        except Exception:
            # 静默忽略所有日志记录错误
            pass
    
    def register_training_thread(self, task_id: str, thread: threading.Thread) -> bool:
        """注册训练线程"""
        with self._data_lock:
            if task_id in self._training_threads:
                self.logger.warning(f"训练线程 {task_id} 已存在，跳过注册")
                return False
            
            self._training_threads[task_id] = thread
            self._active_tasks.add(task_id)
            
            # 初始化任务对应的进程集合
            if task_id not in self._process_by_task:
                self._process_by_task[task_id] = set()
            
            self.logger.info(f"注册训练线程: {task_id}")
            return True
    
    def unregister_training_thread(self, task_id: str) -> bool:
        """注销训练线程，同时清理相关进程"""
        with self._data_lock:
            removed = False
            if task_id in self._training_threads:
                del self._training_threads[task_id]
                removed = True
            
            self._active_tasks.discard(task_id)
            
            # 清理该任务相关的所有注册进程
            if task_id in self._process_by_task:
                process_pids = self._process_by_task[task_id].copy()
                for pid in process_pids:
                    self._cleanup_registered_process(pid)
                del self._process_by_task[task_id]
                self.logger.info(f"清理任务 {task_id} 相关的 {len(process_pids)} 个进程")
            
            if removed:
                self.logger.info(f"注销训练线程: {task_id}")
            return removed
    
    def register_process(self, pid: int, name: str, process_type: str, task_id: Optional[str] = None) -> bool:
        """注册进程到管理器"""
        try:
            # 验证进程是否存在
            psutil.Process(pid)
            
            with self._data_lock:
                process_info = ProcessInfo(
                    pid=pid,
                    name=name,
                    process_type=process_type,
                    start_time=time.time(),
                    task_id=task_id
                )
                
                # 添加到主注册表
                self._registered_processes[pid] = process_info
                
                # 如果有任务ID,添加到任务分组
                if task_id:
                    if task_id not in self._process_by_task:
                        self._process_by_task[task_id] = set()
                    self._process_by_task[task_id].add(pid)
                
                self.logger.info(f"注册进程到管理器: PID={pid}, 类型={process_type}, 任务={task_id}")
                return True
                
        except psutil.NoSuchProcess:
            self.logger.warning(f"尝试注册不存在的进程: PID={pid}")
            return False
    
    def unregister_process(self, pid: int) -> bool:
        """从注册表中移除进程"""
        with self._data_lock:
            if pid in self._registered_processes:
                process_info = self._registered_processes.pop(pid)
                
                # 从任务分组中移除
                task_id = process_info.task_id
                if task_id and task_id in self._process_by_task:
                    self._process_by_task[task_id].discard(pid)
                    if not self._process_by_task[task_id]:  # 如果集合为空，删除键
                        del self._process_by_task[task_id]
                
                self.logger.info(f"从注册表移除进程: PID={pid}, 类型={process_info.process_type}")
                return True
            return False
    
    def _cleanup_registered_process(self, pid: int) -> bool:
        """清理注册表中的指定进程"""
        with self._data_lock:
            if pid not in self._registered_processes:
                return False
            
            process_info = self._registered_processes[pid]
            
            try:
                proc = psutil.Process(pid)
                if proc.is_running():
                    success = self._kill_process_safely(proc, timeout=10)
                    if success:
                        self.logger.info(f"成功终止注册进程: PID={pid}, 类型={process_info.process_type}")
                    else:
                        self.logger.warning(f"无法终止注册进程: PID={pid}")
                    return success
                else:
                    self.logger.info(f"注册进程已停止: PID={pid}")
                    return True
                    
            except psutil.NoSuchProcess:
                self.logger.info(f"注册进程已不存在: PID={pid}")
                return True
            finally:
                # 无论是否成功终止，都从注册表中移除
                self.unregister_process(pid)
    
    def _kill_process_safely(self, process: psutil.Process, timeout: int = 10) -> bool:
        """安全地终止进程"""
        try:
            process_name = process.name()
            process_pid = process.pid
            
            # 首先尝试SIGTERM
            self.logger.info(f"尝试终止进程 {process_pid} ({process_name})")
            process.terminate()
            
            # 等待进程退出
            try:
                process.wait(timeout=timeout)
                self.logger.info(f"进程 {process_pid} ({process_name}) 已退出")
                return True
            except psutil.TimeoutExpired:
                # 如果超时，使用SIGKILL强制终止
                self.logger.warning(f"进程 {process_pid} ({process_name}) 超时，强制终止")
                process.kill()
                process.wait(timeout=5)
                self.logger.info(f"进程 {process_pid} ({process_name}) 已强制终止")
                return True
                
        except psutil.NoSuchProcess:
            self.logger.info(f"进程 {process.pid} 已经不存在")
            return True
        except psutil.AccessDenied:
            self.logger.error(f"没有权限终止进程 {process.pid}")
            return False
        except Exception as e:
            self.logger.error(f"终止进程 {process.pid} 时出错: {e}")
            return False
    
    def stop_training_threads(self, timeout: int = 30) -> Dict[str, bool]:
        """停止所有训练线程"""
        results = {}
        
        with self._data_lock:
            if not self._training_threads:
                return results
            
            self.logger.info(f"正在停止 {len(self._training_threads)} 个训练线程...")
            
            # 设置停止标志
            for task_id in list(self._active_tasks):
                self.logger.info(f"请求停止训练任务: {task_id}")
                results[task_id] = False
        
        # 等待线程结束
        start_time = time.time()
        while time.time() - start_time < timeout:
            with self._data_lock:
                alive_threads = []
                for task_id, thread in list(self._training_threads.items()):
                    if thread.is_alive():
                        alive_threads.append(task_id)
                    else:
                        results[task_id] = True
                        self.unregister_training_thread(task_id)
                
                if not alive_threads:
                    break
            
            time.sleep(1)
        
        # 检查仍在运行的线程
        with self._data_lock:
            for task_id, thread in list(self._training_threads.items()):
                if thread.is_alive():
                    self.logger.warning(f"训练线程 {task_id} 仍在运行")
                    results[task_id] = False
        
        return results
    
    def cleanup_registered_processes(self) -> Dict[int, bool]:
        """清理所有注册的进程"""
        results = {}
        
        with self._data_lock:
            if not self._registered_processes:
                self._safe_log("没有注册的进程需要清理")
                return results
            
            self._safe_log(f"开始清理 {len(self._registered_processes)} 个注册进程...")
            
            # 获取当前注册的所有进程PID
            process_pids = list(self._registered_processes.keys())
        
        # 并发清理进程
        for pid in process_pids:
            success = self._cleanup_registered_process(pid)
            results[pid] = success
        
        self._safe_log(f"进程清理完成，成功: {sum(results.values())}/{len(results)}")
        return results
    
    def graceful_shutdown(self):
        """清理注册表中的进程"""
        if self._shutdown_initiated.is_set():
            return
        
        self._shutdown_initiated.set()
        
        try:
            # 安全记录日志
            self._safe_log("开始关闭进程管理器...")
            
            # 1. 停止训练线程（会自动清理相关进程）
            thread_results = self.stop_training_threads(timeout=20)
            
            # 2. 清理剩余的注册进程
            process_results = self.cleanup_registered_processes()
            
            # 3. 记录完成信息（使用安全日志）
            self._safe_log(f"关闭完成 - 线程: {len(thread_results)}, 进程: {len(process_results)}")
            
            # 4. 最后清理资源（包括关闭日志系统）
            self.cleanup_resources()
            
        except Exception as e:
            self._safe_log(f"关闭时出错: {e}", level="error")
    
    def cleanup_resources(self):
        """清理资源"""
        try:
            # 清理临时文件
            from config import config
            temp_dir = config.paths.temp_dir
            if temp_dir.exists():
                import shutil
                for item in temp_dir.iterdir():
                    try:
                        if item.is_dir():
                            shutil.rmtree(item)
                        else:
                            item.unlink()
                        self.logger.info(f"清理临时文件: {item}")
                    except Exception as e:
                        self.logger.warning(f"清理临时文件失败 {item}: {e}")
            
            # 关闭线程池
            if hasattr(self, '_cleanup_executor'):
                self._cleanup_executor.shutdown(wait=True)
            
            # 清理日志处理器
            logging.shutdown()
            
        except Exception as e:
            self.logger.error(f"清理资源时出错: {e}")
    
    def cleanup_all(self):
        """清理所有资源"""
        if self._shutdown_initiated.is_set():
            return
        
        self._shutdown_initiated.set()
        self._safe_log("开始清理所有资源...")
        
        try:
            # 1. 首先取消所有异步任务
            self._safe_log("取消所有异步任务...")
            async_results = self.cancel_all_async_tasks()
            self._safe_log(f"异步任务取消结果: {len(async_results)} 个任务")
            
            # 2. 停止训练线程
            self._safe_log("停止训练线程...")
            thread_results = self.stop_training_threads()
            self._safe_log(f"训练线程停止结果: {thread_results}")
            
            # 3. 清理注册的进程
            self._safe_log("清理注册的进程...")
            process_results = self.cleanup_registered_processes()
            self._safe_log(f"进程清理结果: {process_results}")
            
            # 4. 关闭线程池
            try:
                self._cleanup_executor.shutdown(wait=True, timeout=5)
            except Exception:
                pass
            
            self._safe_log("资源清理完成")
        except Exception as e:
            self._safe_log(f"清理过程中出现异常: {str(e)}", "error")
    
    def get_registered_processes_info(self) -> Dict[int, ProcessInfo]:
        """获取所有注册进程的信息"""
        with self._data_lock:
            # 更新进程状态
            for pid, process_info in self._registered_processes.items():
                process_info.is_alive()
            return self._registered_processes.copy()
    
    def get_task_processes(self, task_id: str) -> List[ProcessInfo]:
        """获取指定任务的所有进程"""
        with self._data_lock:
            if task_id not in self._process_by_task:
                return []
            
            processes = []
            for pid in self._process_by_task[task_id]:
                if pid in self._registered_processes:
                    processes.append(self._registered_processes[pid])
            return processes
    
    def cleanup_task_processes(self, task_id: str) -> Dict[int, bool]:
        """清理指定任务的所有进程"""
        results = {}
        
        with self._data_lock:
            if task_id not in self._process_by_task:
                self.logger.info(f"任务 {task_id} 没有注册的进程")
                return results
            
            process_pids = self._process_by_task[task_id].copy()
            self.logger.info(f"开始清理任务 {task_id} 的 {len(process_pids)} 个进程")
        
        # 清理进程
        for pid in process_pids:
            success = self._cleanup_registered_process(pid)
            results[pid] = success
        
        return results

    def register_async_task(
        self, 
        task: asyncio.Task, 
        name: str, 
        task_type: str, 
        related_nodes: Optional[List[str]] = None
    ) -> str:
        """注册异步任务到管理器"""
        task_id = f"{task_type}_{int(time.time())}_{uuid.uuid4().hex[:8]}"
        
        with self._data_lock:
            task_info = AsyncTaskInfo(
                task_id=task_id,
                name=name,
                task_type=task_type,
                start_time=time.time(),
                task=task,
                related_nodes=related_nodes or []
            )
            
            # 添加到主注册表
            self._async_tasks[task_id] = task_info
            
            # 按类型分组
            if task_type not in self._async_tasks_by_type:
                self._async_tasks_by_type[task_type] = set()
            self._async_tasks_by_type[task_type].add(task_id)
            
            self.logger.info(f"注册异步任务: {task_id}, 类型={task_type}, 节点={related_nodes}")
            return task_id
    
    def unregister_async_task(self, task_id: str) -> bool:
        """注销异步任务"""
        with self._data_lock:
            if task_id not in self._async_tasks:
                return False
            
            task_info = self._async_tasks[task_id]
            
            # 从类型分组中移除
            if task_info.task_type in self._async_tasks_by_type:
                self._async_tasks_by_type[task_info.task_type].discard(task_id)
                
                # 如果分组为空，删除该分组
                if not self._async_tasks_by_type[task_info.task_type]:
                    del self._async_tasks_by_type[task_info.task_type]
            
            # 从主注册表移除
            del self._async_tasks[task_id]
            
            self.logger.info(f"注销异步任务: {task_id}")
            return True
    
    def cancel_async_tasks_by_type(self, task_type: str, timeout: float = 5.0) -> Dict[str, bool]:
        """根据类型取消异步任务"""
        results = {}
        
        with self._data_lock:
            if task_type not in self._async_tasks_by_type:
                return results
            
            task_ids = self._async_tasks_by_type[task_type].copy()
        
        for task_id in task_ids:
            if task_id in self._async_tasks:
                task_info = self._async_tasks[task_id]
                if task_info.cancel():
                    results[task_id] = True
                    self.logger.info(f"取消异步任务: {task_id}")
                else:
                    results[task_id] = False
                    self.logger.warning(f"无法取消异步任务: {task_id}")
                
                # 注销任务
                self.unregister_async_task(task_id)
        
        return results
    
    def cancel_all_async_tasks(self, timeout: float = 10.0) -> Dict[str, bool]:
        """取消所有异步任务"""
        results = {}
        
        with self._data_lock:
            task_ids = list(self._async_tasks.keys())
        
        for task_id in task_ids:
            if task_id in self._async_tasks:
                task_info = self._async_tasks[task_id]
                if task_info.cancel():
                    results[task_id] = True
                    self.logger.info(f"取消异步任务: {task_id}")
                else:
                    results[task_id] = False
                    self.logger.warning(f"无法取消异步任务: {task_id}")
                
                # 注销任务
                self.unregister_async_task(task_id)
        
        return results
    
    def get_async_tasks_info(self) -> Dict[str, AsyncTaskInfo]:
        """获取所有异步任务信息"""
        with self._data_lock:
            # 清理已完成的任务
            completed_tasks = [
                task_id for task_id, task_info in self._async_tasks.items()
                if task_info.is_done()
            ]
            
            for task_id in completed_tasks:
                self.unregister_async_task(task_id)
            
            return self._async_tasks.copy()
    
    def get_edge_connection_tasks(self) -> List[AsyncTaskInfo]:
        """获取所有边侧连接相关的异步任务"""
        with self._data_lock:
            edge_types = ['edge_connection', 'batch_operation', 'sample_collection', 'health_check']
            edge_tasks = []
            
            for task_type in edge_types:
                if task_type in self._async_tasks_by_type:
                    for task_id in self._async_tasks_by_type[task_type]:
                        if task_id in self._async_tasks:
                            edge_tasks.append(self._async_tasks[task_id])
            
            return edge_tasks

# 全局进程管理器实例
process_manager = ProcessManager()

@contextmanager
def managed_training_task(task_id: str):
    """训练任务上下文管理器 - 自动管理进程注册和清理"""
    thread = threading.current_thread()
    process_manager.register_training_thread(task_id, thread)
    try:
        yield process_manager
    finally:
        process_manager.unregister_training_thread(task_id)

def register_training_process(task_id: str = None):
    """注册训练进程装饰器"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            # 获取当前进程ID
            current_pid = os.getpid()
            
            # 注册进程
            process_manager.register_process(
                current_pid, 
                f"Training-{func.__name__}", 
                "training", 
                task_id
            )
            
            try:
                return func(*args, **kwargs)
            finally:
                # 函数结束时自动注销进程
                process_manager.unregister_process(current_pid)
        return wrapper
    return decorator

def register_yolo_process(task_id: str = None):
    """注册YOLO训练进程装饰器"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            # 获取当前进程ID
            current_pid = os.getpid()
            
            # 注册进程
            process_manager.register_process(
                current_pid, 
                "YOLO Training", 
                "yolo_training", 
                task_id
            )
            
            try:
                return func(*args, **kwargs)
            finally:
                # 函数结束时自动注销进程
                process_manager.unregister_process(current_pid)
        return wrapper
    return decorator

# 异步任务管理装饰器和上下文管理器
@asynccontextmanager
async def managed_async_task(name: str, task_type: str, related_nodes: Optional[List[str]] = None):
    """管理异步任务的上下文管理器"""
    task_id = None
    try:
        # 获取当前运行的任务
        current_task = asyncio.current_task()
        if current_task:
            task_id = process_manager.register_async_task(
                task=current_task,
                name=name,
                task_type=task_type,
                related_nodes=related_nodes
            )
            
        yield task_id
        
    except Exception as e:
        process_manager._safe_log(f"异步任务 {name} 执行异常: {str(e)}", "error")
        raise
    finally:
        if task_id:
            process_manager.unregister_async_task(task_id)

def register_edge_task(task_type: str = "edge_connection"):
    """注册边侧连接任务的装饰器"""
    def decorator(func):
        if asyncio.iscoroutinefunction(func):
            async def async_wrapper(*args, **kwargs):
                # 从参数中提取节点信息
                related_nodes = []
                if 'node_id' in kwargs:
                    related_nodes.append(kwargs['node_id'])
                elif 'node_ids' in kwargs:
                    related_nodes.extend(kwargs['node_ids'])
                elif args and hasattr(args[0], 'node_id'):
                    related_nodes.append(args[0].node_id)
                elif args and isinstance(args[0], list) and args[0] and hasattr(args[0][0], 'node_id'):
                    related_nodes.extend([node.node_id for node in args[0]])
                
                async with managed_async_task(
                    name=f"{func.__name__}",
                    task_type=task_type,
                    related_nodes=related_nodes
                ):
                    return await func(*args, **kwargs)
            return async_wrapper
        else:
            # 对于同步函数，直接返回原函数
            return func
    return decorator

def register_batch_edge_task(task_type: str = "batch_operation"):
    """注册批量边侧操作任务的装饰器"""
    return register_edge_task(task_type)

def register_sample_collection_task():
    """注册样本收集任务的装饰器"""
    return register_edge_task("sample_collection") 